import random

# Pomocne funkcije
def mod_pow(a, n, m):
    result = 1
    a = a % m
    while n > 0:
        if n % 2 == 1:
            result = (result * a) % m
        a = (a * a) % m
        n = n // 2
        
    return result

def miller_rabin(n, k):
    if n <= 3:
        if n == 1:
            return False
        return True
  
    # n prost => n neparan => n = (2 ^ r) * d + 1
    d = n - 1
    r = 0
    while d % 2 == 0:
        r = r + 1
        d = d // 2
        
    for i in range(k):
        a = random.randrange(2, n - 1)
        
        x = mod_pow(a, d, n)

        if x == 1 or x == n - 1:
            continue
            
        wittness = True
        
        for j in range(r - 1):
            x = mod_pow(x, 2, n)
            if x == 1:
                return False
            if x == n - 1: # n - 1 = -1 (mod n)
                wittness = False
                break
        
        if wittness:
            return False
    return True


def get_prime(limit, k = 20):
    is_prime = False
    while not is_prime:
        n = random.randrange(limit)
        is_prime = miller_rabin(n, k)
    return n

# Pomoćna funkcija, prošireni Euklidov algoritam
def gcd(a, b):
    if b == 0:
        return a
    return gcd(b, a % b)

# Pomoćna funkcija, prošireni Euklidov algoritam
def ext_gcd(a, b):
    if b == 0:
        return (a, 1, 0)
    g, x, y = ext_gcd(b, a % b)
    return (g, y, x - a // b * y)

def mod_inv(a, m):
    g, x, y = ext_gcd(a, m)
    if g != 1:
        print("Vrednosti a i m nisu uzajamno proste!")
    else:
        return x % m

class DiffieHellman:
    def __init__(self, g, n):
        self.g = g
        self.n = n
        self.priv = random.randrange(n)
        self.pub = pow(g, self.priv, n)
        
    def calculate_key(self, pub_B):
        self.key = pow(pub_B, self.priv, self.n)

n = get_prime(2**256)
g = random.randrange(2, n)

Adh = DiffieHellman(g, n)
Bdh = DiffieHellman(g, n)

print(f'A priv: {Adh.priv}')
print(f'A pub: {Adh.pub}')

print()

print(f'B priv: {Bdh.priv}')
print(f'B pub: {Bdh.pub}')

print()

print (f'A --[{Adh.pub}]--> B')
print (f'A <--[{Bdh.pub}]-- B')

Adh.calculate_key(Bdh.pub)
Bdh.calculate_key(Adh.pub)

print()

print(f'A: {Adh.key}')
print(f'B: {Bdh.key}')

A priv: 1427868233129745431836084611802856601610057458883993036136196621947682045445
A pub: 219766865127048964507976952846606345896066336617218251052342269829596199650

B priv: 1159408374590420041411236606063861017657718190692398071379499945876929267178
B pub: 985758364139223758749930256396555986789855648896213095670575258654109968029

A --[219766865127048964507976952846606345896066336617218251052342269829596199650]--> B
A <--[985758364139223758749930256396555986789855648896213095670575258654109968029]-- B

A: 2956163243517825289794418419702361851462871897832086085759095087823173333912
B: 2956163243517825289794418419702361851462871897832086085759095087823173333912